#include "ZLibEncryptHelper.h"
#include "zlib.h"
#include <assert.h>
#include <time.h>
#include <string>
using namespace std;

LIB_UTIL_NAMESPACE_BEGIN

#define CHUNK 1024
static byte magic[2] = {0x78, 0x9c};
//static byte trick[2] = {0x98, 0x76};

ZLibEncryptHelper::ZLibEncryptHelper(void)
{
	srand((unsigned)time(NULL));
	trick[0] = (double)rand() / (RAND_MAX + 1) * 0xff;
	trick[1] = (double)rand() / (RAND_MAX + 1) * 0xff;
}

ZLibEncryptHelper::~ZLibEncryptHelper(void)
{
}

int inflate_read(const tDataBuffer& inBuffer, tDataBuffer& outBuffer/*char *source,size_t len,char **dest, size_t* size*/) 
{
	int ret = Z_ERRNO;
	unsigned have;
	z_stream strm;
	byte out[CHUNK];
	byte* dest = NULL;
	int totalsize = 0;

	/* allocate inflate state */
	strm.zalloc = Z_NULL;
	strm.zfree = Z_NULL;
	strm.opaque = Z_NULL;
	strm.avail_in = Z_NULL;
	strm.next_in = Z_NULL;

	//if (inflateInit2(&strm, 47) != Z_OK)
	if (inflateInit(&strm) != Z_OK)
		return ret;

	strm.avail_in = inBuffer.size;
	strm.next_in = (Bytef*)inBuffer.data;

	/* run inflate() on input until output buffer not full */
	do {
		strm.avail_out = CHUNK;
		strm.next_out = out;
		ret = inflate(&strm, Z_NO_FLUSH);
		assert(ret != Z_STREAM_ERROR); /* state not clobbered */
		switch (ret) {
		case Z_NEED_DICT:
			ret = Z_DATA_ERROR; /* and fall through */
		case Z_DATA_ERROR:
		case Z_MEM_ERROR:
			inflateEnd(&strm);
			return ret;
		}
		have = CHUNK - strm.avail_out;
		totalsize += have;
		dest = (byte*)realloc(dest,totalsize);
		memcpy(dest + totalsize - have,out,have);
	} while (strm.avail_out == 0);

	outBuffer.data = new byte[totalsize];
	memcpy(outBuffer.data, dest, totalsize);
	outBuffer.size = totalsize;
	free(dest);

	/* clean up and return */
	(void)inflateEnd(&strm);
	return ret == Z_STREAM_END ? Z_OK : Z_DATA_ERROR;
}

int deflate_write(const tDataBuffer& inBuffer, tDataBuffer& outBuffer/*char* source, size_t len, char** dest, size_t* size*/)
{
	int ret = Z_ERRNO;
	unsigned have;
	z_stream strm;
	byte out[CHUNK];
	byte* dest = NULL;
	int totalsize = 0;


	/* allocate inflate state */
	strm.zalloc = Z_NULL;
	strm.zfree = Z_NULL;
	strm.opaque = Z_NULL;
	strm.avail_in = Z_NULL;
	strm.next_in = Z_NULL;

	if (ret = deflateInit(&strm,Z_BEST_COMPRESSION) != Z_OK)
		return ret;

	strm.avail_in = inBuffer.size;
	strm.next_in = (Bytef*)inBuffer.data;
	int flush = Z_FINISH;
	/* run inflate() on input until output buffer not full */
	do {
		strm.avail_out = CHUNK;
		strm.next_out = out;
		ret = deflate(&strm, flush);
		assert(ret != Z_STREAM_ERROR); /* state not clobbered */
		switch (ret) {
		case Z_NEED_DICT:
			ret = Z_DATA_ERROR; /* and fall through */
		case Z_DATA_ERROR:
		case Z_MEM_ERROR:
			inflateEnd(&strm);
			return ret;
		}
		have = CHUNK - strm.avail_out;
		totalsize += have;
		dest = (byte*)realloc(dest,totalsize);
		memcpy(dest + totalsize - have,out,have);
	} while (strm.avail_out == 0);
	
	outBuffer.data = new byte[totalsize];
	memcpy(outBuffer.data, dest, totalsize);
	outBuffer.size = totalsize;
	free(dest);

	/* clean up and return */
	(void)deflateEnd(&strm);
	return ret == Z_STREAM_END ? Z_OK : Z_DATA_ERROR;
}

const std::string base64_chars = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";

inline bool is_base64(unsigned char c) {
	return (isalnum(c) || (c == '+') || (c == '/'));
}

static const char base64_table[] =
{ 'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M',
'N', 'O', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z',
'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm',
'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z',
'0', '1', '2', '3', '4', '5', '6', '7', '8', '9', '+', '/', '\0'
};

static const char base64_pad = '=';

static const short base64_reverse_table[256] = {
	-2, -2, -2, -2, -2, -2, -2, -2, -2, -1, -1, -2, -2, -1, -2, -2,
	-2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2,
	-1, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, 62, -2, -2, -2, 63,
	52, 53, 54, 55, 56, 57, 58, 59, 60, 61, -2, -2, -2, -2, -2, -2,
	-2,  0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14,
	15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, -2, -2, -2, -2, -2,
	-2, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40,
	41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, -2, -2, -2, -2, -2,
	-2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2,
	-2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2,
	-2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2,
	-2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2,
	-2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2,
	-2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2,
	-2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2,
	-2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2, -2
};

void base64_encode(const tDataBuffer& inBuffer, tDataBuffer& outBuffer)
{
	const unsigned char *current = inBuffer.data;
	unsigned char *p;
	unsigned char *result;
	int length = inBuffer.size;
	result = (unsigned char *)malloc(((length + 2) / 3) * 4 * sizeof(char));
	p = result;

	while (length > 2) { /* keep going until we have less than 24 bits */
		*p++ = base64_table[current[0] >> 2];
		*p++ = base64_table[((current[0] & 0x03) << 4) + (current[1] >> 4)];
		*p++ = base64_table[((current[1] & 0x0f) << 2) + (current[2] >> 6)];
		*p++ = base64_table[current[2] & 0x3f];

		current += 3;
		length -= 3; /* we just handle 3 octets of data */
	}

	/* now deal with the tail end of things */
	if (length != 0) {
		*p++ = base64_table[current[0] >> 2];
		if (length > 1) {
			*p++ = base64_table[((current[0] & 0x03) << 4) + (current[1] >> 4)];
			*p++ = base64_table[(current[1] & 0x0f) << 2];
			*p++ = base64_pad;
		} else {
			*p++ = base64_table[(current[0] & 0x03) << 4];
			*p++ = base64_pad;
			*p++ = base64_pad;
		}
	}
	outBuffer.size = p - result;
	outBuffer.data = new byte[outBuffer.size + 1];
	memcpy(outBuffer.data, result, outBuffer.size);
	outBuffer.data[outBuffer.size] = 0;
	free(result);
}

//void base64_encode(const tDataBuffer& inBuffer, tDataBuffer& outBuffer/*unsigned char const* bytes_to_encode, unsigned int in_len*/) {
//	/*string ret;
//	int i = 0;
//	int j = 0;
//	unsigned char char_array_3[3];
//	unsigned char char_array_4[4];
//
//	size_t in_len = inBuffer.size;
//	byte* bytes_to_encode = inBuffer.data;
//	while (in_len--) {
//		char_array_3[i++] = *(bytes_to_encode++);
//		if (i == 3) {
//			char_array_4[0] = (char_array_3[0] & 0xfc) >> 2;
//			char_array_4[1] = ((char_array_3[0] & 0x03) << 4) + ((char_array_3[1] & 0xf0) >> 4);
//			char_array_4[2] = ((char_array_3[1] & 0x0f) << 2) + ((char_array_3[2] & 0xc0) >> 6);
//			char_array_4[3] = char_array_3[2] & 0x3f;
//
//			for(i = 0; (i <4) ; i++)
//				ret += base64_chars[char_array_4[i]];
//			i = 0;
//		}
//	}
//
//	if (i)
//	{
//		for(j = i; j < 3; j++)
//			char_array_3[j] = '\0';
//
//		char_array_4[0] = (char_array_3[0] & 0xfc) >> 2;
//		char_array_4[1] = ((char_array_3[0] & 0x03) << 4) + ((char_array_3[1] & 0xf0) >> 4);
//		char_array_4[2] = ((char_array_3[1] & 0x0f) << 2) + ((char_array_3[2] & 0xc0) >> 6);
//		char_array_4[3] = char_array_3[2] & 0x3f;
//
//		for (j = 0; (j < i + 1); j++)
//			ret += base64_chars[char_array_4[j]];
//
//		while((i++ < 3))
//			ret += '=';
//
//	}
//	byte* data = new byte[ret.length() + 1];
//	memcpy(data, ret.c_str(), ret.length() + 1);
//	outBuffer.data = data;
//	outBuffer.size = ret.length() + 1;*/
//	EVP_ENCODE_CTX ctx;
//	EVP_EncodeInit(&ctx);
//	byte* buffer = new byte[inBuffer.size];
//	int size = 0;
//	int out = 0;
//	EVP_EncodeUpdate(&ctx, buffer, &out, inBuffer.data, inBuffer.size);
//	size += out;
//	EVP_EncodeFinal(&ctx, buffer + out, &out);
//	size += out;
//	outBuffer.data = new byte[size + 1];
//	outBuffer.size = size;
//	memcpy(buffer, outBuffer.data, size);
//	delete []buffer;
//	outBuffer.data[size] = 0;
//}

void base64_decode(const tDataBuffer& inBuffer, tDataBuffer& outBuffer)
{
	const unsigned char *current = inBuffer.data;
	int ch, i = 0, j = 0, k;
	/* this sucks for threaded environments */
	unsigned char *result;
	int length = inBuffer.size;
	result = (unsigned char *)malloc(length + 1);

	/* run through the whole string, converting as we go */
	while ((ch = *current++) != '\0' && length-- > 0) {
		if (ch == base64_pad) break;

		ch = base64_reverse_table[ch];
		if (ch < 0)  /* a space or some other separator character, we simply skip over */
			continue;

		switch(i % 4) {
		case 0:
			result[j] = ch << 2;
			break;
		case 1:
			result[j++] |= ch >> 4;
			result[j] = (ch & 0x0f) << 4;
			break;
		case 2:
			result[j++] |= ch >>2;
			result[j] = (ch & 0x03) << 6;
			break;
		case 3:
			result[j++] |= ch;
			break;
		}
		i++;
	}

	outBuffer.size = j;
	outBuffer.data = new byte[j + 1];
	memcpy(outBuffer.data, result, j);
	outBuffer.data[j] = 0;
	free(result);
}

//void base64_decode(const tDataBuffer& inBuffer, tDataBuffer& outBuffer/*string const& encoded_string*/) {
//	EVP_ENCODE_CTX ctx;
//	EVP_DecodeInit(&ctx);
//	byte* buffer = new byte[inBuffer.size * 2];
//	int size = 0;
//	int out = 0;
//	int shift = 0;
//	do {
//		EVP_DecodeUpdate(&ctx, buffer + ctx.num, &out, inBuffer.data + shift, inBuffer.size - shift);
//		shift += out;
//		size += ctx.num;
//	} while(out);
//	
//	size += out;
//	EVP_DecodeFinal(&ctx, buffer + out, &out);
//	size += out;
//	outBuffer.data = new byte[size];
//	outBuffer.size = size;
//	memcpy(buffer, outBuffer.data, size);
//	delete []buffer;
//	/*size_t in_len = inBuffer.size;
//	byte* encoded_string = inBuffer.data;
//	int i = 0;
//	int j = 0;
//	int in_ = 0;
//	unsigned char char_array_4[4], char_array_3[3];
//	string ret;
//
//	while (in_len-- && ( encoded_string[in_] != '=') && is_base64(encoded_string[in_])) {
//		char_array_4[i++] = encoded_string[in_]; in_++;
//		if (i ==4) {
//			for (i = 0; i <4; i++)
//				char_array_4[i] = base64_chars.find(char_array_4[i]);
//
//			char_array_3[0] = (char_array_4[0] << 2) + ((char_array_4[1] & 0x30) >> 4);
//			char_array_3[1] = ((char_array_4[1] & 0xf) << 4) + ((char_array_4[2] & 0x3c) >> 2);
//			char_array_3[2] = ((char_array_4[2] & 0x3) << 6) + char_array_4[3];
//
//			for (i = 0; (i < 3); i++)
//				ret += char_array_3[i];
//			i = 0;
//		}
//	}
//
//	if (i) {
//		for (j = i; j <4; j++)
//			char_array_4[j] = 0;
//
//		for (j = 0; j <4; j++)
//			char_array_4[j] = base64_chars.find(char_array_4[j]);
//
//		char_array_3[0] = (char_array_4[0] << 2) + ((char_array_4[1] & 0x30) >> 4);
//		char_array_3[1] = ((char_array_4[1] & 0xf) << 4) + ((char_array_4[2] & 0x3c) >> 2);
//		char_array_3[2] = ((char_array_4[2] & 0x3) << 6) + char_array_4[3];
//
//		for (j = 0; (j < i - 1); j++) ret += char_array_3[j];
//	}
//	byte* data = new byte[ret.length()];
//	memcpy(data, ret.c_str(), ret.length());
//	outBuffer.data = data;
//	outBuffer.size = ret.length();*/
//}

bool ZLibEncryptHelper::DecryptData(const tDataBuffer& inBuffer, tDataBuffer& outBuffer, EncryptType type)
{
	switch(type)
	{
	case Binary:
		inBuffer.data[0] = magic[0];
		inBuffer.data[1] = magic[1];
		//return inflate_read((char*)inBuffer.data, inBuffer.size, (char**)&outBuffer.data, &outBuffer.size) == 0;
		return inflate_read(inBuffer, outBuffer) == 0;
	case BinaryNonEncrypt:
		outBuffer = inBuffer;
		return true;
	case Base64:
		{
			tDataBuffer buffer;
			base64_decode(inBuffer, buffer);
			buffer.data[0] = magic[0];
			buffer.data[1] = magic[1];
			//return inflate_read((char*)buffer.data, buffer.size, (char**)&outBuffer.data, &outBuffer.size) == 0;
			return inflate_read(buffer, outBuffer) == 0;
		}
	case Base64NoeEncrypt:
		base64_decode(inBuffer, outBuffer);
		return true;
	default:
		return false;
	}
}

bool ZLibEncryptHelper::EncryptData(const tDataBuffer& inBuffer, tDataBuffer& outBuffer, EncryptType type)
{
	switch(type)
	{
	case Binary:
		{
			//bool ret = deflate_write((char*)inBuffer.data, inBuffer.size, (char**)&outBuffer.data, &outBuffer.size) == 0;
			bool ret = deflate_write(inBuffer, outBuffer) == 0;
			outBuffer.data[0] = trick[0];
			outBuffer.data[1] = trick[1];
			trick[0] = outBuffer.data[outBuffer.size >> 1];
			trick[1] = outBuffer.data[outBuffer.size - 2];
			return ret;
		}
	case BinaryNonEncrypt:
		outBuffer = inBuffer;
		return true;
#ifdef _WINDOWS
	case Base64:
		{
			tDataBuffer encryptBuffer;
			if(!EncryptData(inBuffer, encryptBuffer, Binary))
				return false;
			base64_encode(encryptBuffer, outBuffer);
			return true;
		}
	case Base64NoeEncrypt:
		base64_encode(inBuffer, outBuffer);
		return true;
#endif //_WINDOWS
	default:
		return false;
	}
}

EncryptHelper* EncryptHelper::CreateEncryptHelper()
{
	return new ZLibEncryptHelper;
}

void EncryptHelper::ReleaseEncryptHelper(EncryptHelper* helper)
{
	if(helper == NULL)
		return;
	delete (ZLibEncryptHelper*)helper;
}

LIB_UTIL_NAMESPACE_END
